import numpy as np
import scipy.special
from matplotlib import pyplot as plt
import random

################################################################
# 哈里鹰优化算法
################################################################


class HHO():
    """哈里鹰优化算法"""

    def __init__(self, FitFunction, LB, UB, dim, population_size, b, max_iter):
        """初始化物种优化算法的参数"""
        self.LB = LB  # 搜索空间的左边界
        self.UB = UB  # 搜索空间的右边界
        self.dim = dim  # 位置信息的维度
        self.population_size = population_size  # 种群大小
        self.max_iter = max_iter  # 迭代次数
        self.b = b  # 螺旋式更新位置的参数
        self.X = np.zeros((population_size, dim))  # 种群的位置信息
        self.FitFunction = FitFunction
        self.gBest_score = np.inf  # 最优解对应的适应度值
        self.gBest_curve = []  # 每轮的最优解的适应度值
        self.gBest_X = np.zeros(dim)  # 每次迭代得到的最佳位置

    def Initialize_population(self):
        """initialize population 初始化种群数据"""
        # 初始化种群数据X
        for i in range(self.population_size):
            for j in range(self.dim):
                self.X[i, j] = (self.UB[j] - self.LB[j]) * \
                    np.random.random() + self.LB[j]

    def BorderCheck_UpdateFitness(self):
        """边界检查与更新适应度值"""
        for i in range(self.population_size):
            for j in range(self.dim):
                # 检查边界
                for j in range(self.dim):
                    if (self.X[i, j] > self.UB[j]):
                        # 不符合上边界
                        self.X[i, j] = self.UB[j]

                    if (self.X[i, j] < self.LB[j]):
                        # 不符合下边界
                        self.X[i, j] = self.LB[j]
                # 计算所有物种的fitness
                fitness = self.CalculateFitness(self.X[i, :])  # 获取物种的适应度值
                # 更新最优解的适应度值和物种的最优解的位置信息
                if fitness < self.gBest_score:
                    self.gBest_score = fitness
                    self.gBest_X = self.X[i, :].copy()

    def CalculateFitness(self, X):
        """自定义计算适应度函数,计算单个物种的适应度值"""
        fitness = self.FitFunction(X)
        return fitness

    def Levy(self):
        """Levy飞行机制"""
        beta = 1.5
        sigma = (scipy.special.gamma(1+beta)*np.sin(np.pi*beta/2) /
                 (scipy.special.gamma((1+beta)/2)*beta**((beta-1)/2)))**(1/beta)
        u = np.random.random(size=(1, self.dim))*sigma
        v = np.random.random(size=(1, self.dim))
        output = u/abs(v) ** (1/beta)
        return output

    def Optimize(self):
        self.Initialize_population()
        t = 0
        iter = 0
        while iter < self.max_iter:
            # 固定工作
            for i in range(self.population_size):
                for j in range(self.dim):
                    # 检查边界
                    for j in range(self.dim):
                        if (self.X[i, j] > self.UB[j]):
                            # 不符合上边界
                            self.X[i, j] = self.UB[j]

                        if (self.X[i, j] < self.LB[j]):
                            # 不符合下边界
                            self.X[i, j] = self.LB[j]
                    # 计算所有物种的fitness
                    fitness = self.CalculateFitness(self.X[i, :])  # 获取物种的适应度值
                    iter = iter+1
                    # 更新最优解的适应度值和物种的最优解的位置信息
                    if fitness < self.gBest_score:
                        self.gBest_score = fitness
                        self.gBest_X = self.X[i, :].copy()

            E_1 = 2*(1-iter/self.max_iter)
            # 更新物种的位置信息
            for i in range(self.population_size):
                # 自定义三个过程的所需参数
                E_0 = 2 * np.random.rand()-1  # -1<E0<1
                # 猎物逃跑的能量
                Escaping_Energy = E_0*E_1
                if abs(Escaping_Energy) >= 1:
                    q = np.random.rand()
                    # 寻找猎物,搜寻猎物
                    rand_index = np.random.randint(
                        low=0, high=self.population_size)
                    X_rand = self.X[rand_index, :]
                    if q < 0.5:
                        self.X[i, :] = X_rand-np.random.rand() * \
                            abs(X_rand-2*np.random.rand()*self.X[i, :])
                    else:
                        self.X[i, :] = (self.gBest_X - np.mean(self.X))-np.random.rand()*(
                            (self.UB-self.LB)*np.random.rand()+self.LB)
                else:
                    r = np.random.rand()

                    if r > 0.5 and abs(Escaping_Energy) < 0.5:
                        self.X[i, :] = self.gBest_X-Escaping_Energy * \
                            abs(self.gBest_X-self.X[i, :])
                    if r > 0.5 and abs(Escaping_Energy) > 0.5:
                        Jump_strength = 2*(1-np.random.rand())  # 随机步长
                        self.X[i, :] = (self.gBest_X-self.X[i, :])-Escaping_Energy * \
                            abs(Jump_strength*self.gBest_X-self.X[i, :])
                    if r < 0.5 and abs(Escaping_Energy) > 0.5:
                        Jump_strength = 2*(1-np.random.rand())  # 随机步长
                        X1 = self.gBest_X-Escaping_Energy * \
                            abs(Jump_strength*self.gBest_X-self.X[i, :])
                        iter = iter+2
                        if self.FitFunction(X1) < self.FitFunction(self.X[i, :]):
                            self.X[i, :] = X1
                        else:
                            X2 = self.gBest_X-Escaping_Energy * \
                                abs(Jump_strength*self.gBest_X-self.X[i, :]) + np.random.random(
                                    size=(1, self.dim))*self.Levy()
                            X2 = X2.reshape(self.dim)
                            iter = iter+1
                            if self.FitFunction(X2) < self.FitFunction(self.X[i, :]):
                                self.X[i, :] = X2

                    if r < 0.5 and abs(Escaping_Energy) < 0.5:
                        Jump_strength = 2*(1-np.random.rand())  # 随机步长
                        X1 = self.gBest_X-Escaping_Energy * \
                            abs(Jump_strength*self.gBest_X-np.mean(self.X))

                        iter = iter+2

                        if self.FitFunction(X1) < self.FitFunction(self.X[i, :]):
                            self.X[i, :] = X1
                        else:
                            X2 = self.gBest_X-Escaping_Energy*abs(Jump_strength*self.gBest_X-np.mean(
                                self.X)) + np.random.random(size=(1, self.dim))*self.Levy()
                            X2 = X2.reshape(self.dim)
                            iter = iter+1
                            if self.FitFunction(X2) < self.FitFunction(self.X[i, :]):
                                self.X[i, :] = X2

            # 更新每轮的最优解的适应度值
            self.gBest_curve.append(self.gBest_score)
            if (t % 10 == 0):
                print('第{}轮训练完成!, 该轮最优解的适应度值为: {}'.format(
                    t, self.gBest_curve[t]))
            t += 1

        return self.gBest_curve, self.gBest_X


# if __name__ == '__main__':
#     whale_nums = 30
#     dim = 30
#     max_iter = dim*10000
#     LB = -10 * np.ones(dim)
#     UB = 10 * np.ones(dim)
#     Func = Fitness(min=LB, max=UB, dim=dim)
#     FitFunction = Func.F1
#     b = 2
#     Gbest_curve, GbestX = HHO(FitFunction, LB=LB, UB=UB,
#                               dim=dim, b=b, population_size=whale_nums, max_iter=max_iter).Optimize()
#     # 可视化训练图像
#     plt.figure(1)
#     plt.plot(Gbest_curve, 'r-', linewidth=2)
#     plt.xlabel('Iterations')
#     plt.ylabel('Fitness')
#     plt.xlim((0, len(Gbest_curve)))
#     plt.grid()
#     plt.title('WOA', fontsize='large')
#     plt.show()
